#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
@author: xjx
@time: 2024/4/7 15:53 
@file: load_analysis_2.py
@project: nanchangproject
@describe: 用于载荷分析的类
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt



class LoadModel(object):

    # 实例变量
    def __init__(self, ig, torque_max, windward_area, air_density, air_coefficient,
                 road_coefficient, tire_radius, ratio_r, ratio_1, ratio_2, ratio_3,
                 ratio_4, ratio_5, ratio_6):
        self.ig = ig  # 主减速比
        self.torque_max = torque_max  # 最大扭矩
        self.windward_area = windward_area  # 车辆迎风面积
        self.air_density = air_density  # 空气密度
        self.air_coefficient = air_coefficient  # 空气阻力系数
        self.road_coefficient = road_coefficient  # 滚阻系数
        self.tire_radius = tire_radius  # 轮胎滚动半径
        self.gear_ratios = [ratio_r, ratio_1, ratio_2, ratio_3, ratio_4, ratio_5, ratio_6]  # 变速器档速比 R123456
        self.std_torque = self.torque_max / 5
        self.gravity = 9.81  # 重力加速度 m/s^2
        self.std_v = 0.1  # 速度噪声标准差
        self.std_inverse_m = 100  # 载荷噪声标准差(非倒数)
        self.std_alpha = 0.03  # 坡度噪声标准差

    @staticmethod
    def transmission_efficiency(u_k):
        """
        用于确定当前传动效率
        @param u_k: 输入扭矩
        @return: 当前传动效率
        """
        torque = [0.0653, 0.0816, 0.1, 0.1490, 0.1939, 0.2490, 0.3, 0.3816, 0.4837, 0.4959, 0.6061, 0.6224, 0.7735,
                  0.7939, 1]  # 扭矩百分比 共15档
        eta = [80.2105, 84.3158, 87.3684, 90.9474, 92.3158, 93.4737, 94.2105, 94.7368, 95.3684, 95.4737, 95.7895,
               95.8947, 96.3158, 96.4211, 96.3158]  # 传动效率
        closest_torque = min(torque, key=lambda x: abs(x - u_k))  # 查找最接近的扭矩
        eta_position = torque.index(closest_torque)  # 确定传动效率位置
        eta = eta[eta_position]  # 确定传动效率
        return eta

    def determine_gear(self, engine_rpm, v_kmph):
        """
        @param engine_rpm: 发动机转速
        @param v_kmph: 车速，单位kmph
        @return: 挡位，档速比
        """
        gear_ratios = self.gear_ratios  # 变速器档速比 R123456
        v_mpm = v_kmph * 1000 / 60  # 车速单位转换
        if v_mpm != 0:
            actual_it = (engine_rpm / v_mpm)  # 计算实际传动比
        else:
            actual_it = 0
        closest_gear = min(gear_ratios, key=lambda x: abs(x - actual_it))  # 查找最接近的已知档速比
        gear_position = gear_ratios.index(closest_gear)  # 确定挡位
        it = gear_ratios[gear_position]  # 确定档速比
        return gear_position, it

    def generate_noise(self):
        """
        用于生成噪声
        @return:
        """
        omega_Torque = np.random.normal(0, self.std_torque)  # 扭矩噪声
        omega_v = np.random.normal(0, self.std_v)  # 速度噪声
        omega_inverse_m = np.random.normal(0, self.std_inverse_m)  # 载荷噪声
        omega_alpha = np.random.normal(0, self.std_alpha)  # 坡度噪声
        return omega_Torque, omega_v, omega_inverse_m, omega_alpha

    def forward_euler_discretization(self, x_k, u_k, v_mps, Ts, it, eta, omega_Torque, omega_inverse_m,
                                     omega_alpha, delta):
        """
        用于前向欧拉离散化
        @param x_k:
        @param u_k:
        @param v_mps:
        @param Ts:
        @param it:
        @param eta:
        @param omega_Torque:
        @param omega_inverse_m:
        @param omega_alpha:
        @param delta:
        @return:
        """
        Torque_hat = max(0, min(1, u_k * self.torque_max + omega_Torque))  # 考虑噪声的扭矩
        v_hat = max(0, min(90.0 / 3.6, v_mps))  # 不考虑噪声的速度
        inverse_m_hat = max(1 / 6100, min(1 / 2800, 1 / (1 / x_k[1] + omega_inverse_m)))  # 考虑噪声的载荷
        alpha_hat = max(-15, min(15, x_k[2] + omega_alpha))  # 考虑噪声的坡度
        x_kplus1 = [0.0, 0.0, 0.0]
        x_kplus1[0] = max(0, min(90 / 3.6, v_hat + Ts / delta * (
                it * self.ig * eta / self.tire_radius * Torque_hat * inverse_m_hat \
                - self.air_density * self.air_coefficient * self.windward_area / 2 * inverse_m_hat * v_hat ** 2 \
                - self.road_coefficient * self.gravity * np.cos(alpha_hat) - self.gravity * np.sin(alpha_hat))))
        x_kplus1[1] = inverse_m_hat
        x_kplus1[2] = alpha_hat
        return x_kplus1

    def extended_average_filter(self, x_k, u_k, P, x_pred, v_mps, Ts, it, eta, omega_v, delta):  # 扩展均值滤波
        Q = np.diag([self.std_v ** 2, self.std_inverse_m ** 2, self.std_alpha ** 2])  # 状态噪声协方差矩阵
        R = np.diag([self.std_v ** 2])  # 观测噪声协方差矩阵
        H = np.array([[1, 0, 0]])
        V = np.array([[1]])

        a11 = 1 - Ts * self.air_density * self.air_coefficient * self.windward_area / delta * x_k[1] * x_k[0]
        a12 = it * self.ig * eta / self.tire_radius / delta * Ts * u_k * self.torque_max - self.air_density * self.air_coefficient * self.windward_area / 2 / delta * Ts * \
              x_k[0] ** 2
        a13 = Ts * self.road_coefficient * self.gravity / delta * np.sin(x_k[2]) - Ts * self.gravity / delta * np.cos(
            x_k[2])
        A = np.array([[a11, a12, a13], [0, 1, 0], [0, 0, 1]], dtype=float)
        w11 = 1 - Ts * self.air_density * self.air_coefficient * self.windward_area / delta * x_k[1] * x_k[0]
        w12 = it * self.ig * eta / self.tire_radius / delta * Ts * u_k * self.torque_max
        w13 = -Ts * self.gravity / delta
        W = np.array([[w11, w12, w13], [0, 1, 0], [0, 0, 1]], dtype=float)
        z = H @ x_pred + V @ np.array([omega_v])  # 预测下一观测值
        z[0] = max(0, min(90 / 3.6, z[0]))  # 观测值限幅
        P_pred = A @ P @ A.T + W @ Q @ W.T
        # P_pred = A@P@A.T+Q
        K = P_pred @ H.T @ np.linalg.inv(H @ P_pred @ H.T + R)
        # x = x_pred+K@(z-H@x_pred)
        x = x_pred + K @ (v_mps - H @ x_pred)  # 用真实车速替代z
        P = (np.eye(3) - K @ H) @ P_pred
        return x, P

    def main_loop_actual(self, time, engine_rpm, v_kmph, v_mps, Torque):
        Ts = 1  # 时间步长 s
        u0 = 350  # 初始控制变量 扭矩
        x0 = [40 / 3.6, 1 / 4500, 0]  # 初始状态变量 [速度,1/载荷,坡度]
        P0 = np.eye(3)  # 初始状态估计协方差矩阵
        x_k = x0
        P_k = P0
        x_k_dict = {}  # 用于存储每步x_k的字典
        try:
            for i in range(len(engine_rpm)):
                if i != len(engine_rpm) - 1:
                    Ts = 1  # time[i + 1] - time[i]
                omega_Torque, omega_v, omega_inverse_m, omega_alpha = self.generate_noise()
                eta = self.transmission_efficiency(Torque[i])
                current_gear, it = self.determine_gear(engine_rpm[i], v_kmph[i])
                delta = 1 + 0.04 * it + 0.0025 * it * self.ig  # 旋转质量换算系数
                x_pred = self.forward_euler_discretization(
                    x_k=x_k,
                    u_k=Torque[i],
                    v_mps=v_mps[i],
                    Ts=Ts,
                    it=it,
                    eta=eta,
                    omega_Torque=omega_Torque,
                    omega_inverse_m=omega_inverse_m,
                    omega_alpha=omega_alpha,
                    delta=delta
                )  # 预测下一状态
                x_k, P_k = self.extended_average_filter(
                    x_k=x_k,
                    u_k=Torque[i],
                    P=P_k,
                    x_pred=x_pred,
                    v_mps=v_mps[i],
                    Ts=Ts,
                    it=it,
                    eta=eta,
                    omega_v=omega_v,
                    delta=delta
                )
                x_k_dict[i] = [x_pred[0], 1 / x_pred[1], x_pred[2]]  # 保存当前x_k
        except Exception as e:
            print("发生异常：", str(e))
            pass
        return x_k_dict, P_k  # 返回状态变量

    def plot_data(self, data_dict, picture_path: str):
        """
        用于绘制载荷预测图
        @param data_dict: 一个向量矩阵，包含载荷预测值
        @param picture_path: 存储图片的路径
        @return:
        """
        start = 0
        length = 5000
        time_steps = list(data_dict.keys())[start:start + length]  # 仅绘制前边的数据
        reciprocals = [data_dict[i][1] for i in time_steps]
        plt.figure(figsize=(10, 6))
        plt.plot(time_steps, reciprocals, label='Mass')
        plt.xlabel('Time Step (s)')
        plt.ylabel('Load Prediction (kg)')
        my_title = picture_path.split('/')[-1]
        plt.title(my_title)
        plt.grid(True)
        plt.savefig(picture_path)


    @staticmethod
    def read_csv_new(csv_file: str):
        """
        由于日方给的数据文件和实验的csv文件不太一样所以需要另外处理
        @param csv_file: csv文件路径
        @return:
        """
        time = []
        engine_rpm = []
        v_kmph = []
        v_mps = []
        Torque = []
        df = pd.read_csv(csv_file)
        df = df[df['发动机转速（RPM）'] != 0]
        df['采集时间'] = pd.to_datetime(df['采集时间'])
        selected_columns = [
            '采集时间',
            '发动机转速（RPM）',
            '车速（km/h）',
            '发动机净输出扭矩或发动机实际扭矩/指示扭矩（%）',
            '摩擦扭矩（%）'
        ]
        df_selected = df[selected_columns]
        df_sorted = df_selected.sort_values(by='采集时间')
        time_column = df_sorted['采集时间']
        engine_rpm_column = df_sorted['发动机转速（RPM）']
        speed_column = df_sorted['车速（km/h）']
        engine_torque_column = df_sorted['发动机净输出扭矩或发动机实际扭矩/指示扭矩（%）']
        friction_torque_column = df_sorted['摩擦扭矩（%）']
        for i in range(len(engine_rpm_column)):
            try:
                rpm_value = float(engine_rpm_column[i])
                speed_value = float(speed_column[i])
                engine_torque_value = float(engine_torque_column[i])
                friction_torque_value = float(friction_torque_column[i])  # 处理摩擦扭矩列
                if rpm_value != 0 and speed_value != 0:
                    time.append(time_column[i])
                    engine_rpm.append(rpm_value)
                    v_kmph.append(speed_value)
                    speed_value_mps = speed_value / 3.6  # 车速转换为 m/s
                    v_mps.append(speed_value_mps)
                    total_torque = (engine_torque_value + friction_torque_value) / 100  # 相加得到最终发动机扭矩
                    Torque.append(total_torque)
            except Exception as e:
                # print('读取csv文件发生错误', e)
                continue
        return time, engine_rpm, v_kmph, v_mps, Torque

    @staticmethod
    def read_csv_actual(csv_file):
        """
        用于读取csv文件
        @param csv_file: csv文件路径
        @return:
        """
        time = []
        engine_rpm = []
        v_kmph = []
        v_mps = []
        Torque = []
        df = pd.read_csv(csv_file)
        time_column = df['时间戳']
        engine_rpm_column = df['发动机转速']
        speed_column = df['车速']
        engine_torque_column = df['发动机净输出扭矩或发动机实际扭矩/指示扭矩']
        friction_torque_column = df['摩擦扭矩']
        for i in range(len(engine_rpm_column)):
            try:
                rpm_value = float(engine_rpm_column[i])
                speed_value = float(speed_column[i])
                engine_torque_value = float(engine_torque_column[i])
                friction_torque_value = float(friction_torque_column[i])  # 处理摩擦扭矩列
                if rpm_value != 0 and speed_value != 0:
                    time.append(time_column[i])
                    engine_rpm.append(rpm_value)
                    v_kmph.append(speed_value)
                    speed_value_mps = speed_value / 3.6  # 车速转换为 m/s
                    v_mps.append(speed_value_mps)
                    total_torque = (engine_torque_value + friction_torque_value) / 100  # 相加得到最终发动机扭矩
                    Torque.append(total_torque)
            except Exception as e:
                print(e)
                continue
        return time, engine_rpm, v_kmph, v_mps, Torque


if __name__ == '__main__':
    my_load = LoadModel(ig=4.11,
                        torque_max=450,
                        windward_area=2.29 * 2.16,
                        air_density=1.225,
                        air_coefficient=0.3,
                        road_coefficient=0.015,
                        tire_radius=0.34284,
                        ratio_r=5.701,
                        ratio_1=5.797,
                        ratio_2=3.434,
                        ratio_3=1.862,
                        ratio_4=1.29,
                        ratio_5=1.0,
                        ratio_6=0.759)
    my_file_name = 'vehicle_data_combined_1205-667(1)-1.csv'
    picture_name = 'myplot3.png'
    time, engine_rpm, v_kmph, v_mps, Torque = my_load.read_csv_actual(my_file_name)
    state_dict, final_covariance = my_load.main_loop_actual(time, engine_rpm, v_kmph, v_mps, Torque)  # 运行主循环
    my_load.plot_data(state_dict, picture_name)
